""""""
import os
import enum
import torch
from PIL import Image
from .util import TencentTranslate, ContainsAnyDict


class Color(enum.StrEnum):
    Red = 'red'
    Green = 'green'
    White = 'white'
    Default = 'white'


class DataType(enum.StrEnum):
    Tensor = 'tensor'
    TensorData = 'tensor_data'
    Image = 'image'
    Json = 'json'
    String = 'string'
    Int = 'int'
    Float = 'float'
    List = 'list'
    Unknown = 'unknown'


class DisplayInfo:
    """
    信息展示
    """
    
    CATEGORY = "destiny/工具"

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "input_data": ("*", {}),
                "size": ("INT", {"default": 12}),
                "color": ([Color.Default, Color.Red, Color.Green, Color.White], {"default": Color.Default}),
            }
        }
    
    @classmethod
    def VALIDATE_INPUTS(s, input_types):
        return True

    RETURN_TYPES = ()
    RETURN_NAMES = ()
    OUTPUT_NODE = True

    FUNCTION = "execute"

    def execute(self, input_data, size, color):
        data_type = self._detect_type(input_data)
        text = [f"类型: {data_type}"]
        # print(text)
        if data_type == DataType.List:
            if len(input_data) > 0:
                data_type = self._detect_type(input_data[0])
                if data_type == DataType.String or data_type == DataType.Int or data_type == DataType.Float:
                    input_data = ",".join([f"{item}" for item in input_data])
                else:
                    text[0] = f"类型: {DataType.List}-{data_type}"
                    input_data = input_data[0]
        if data_type == DataType.String or data_type == DataType.Int or data_type == DataType.Float:
            text.append(f"值: {input_data}")
        elif data_type == DataType.Tensor:
            text.append(f"形状: {input_data.size()}")
        return {
            "ui": {
                "text": text,
                "size": [size],
                'color': [color]
            },
            "result": ()
        }
    
    def _detect_type(self, data):
        if isinstance(data, torch.Tensor):
            return DataType.Tensor if data.dim() > 2 else DataType.TensorData
        elif isinstance(data, Image.Image):
            return DataType.Image
        elif isinstance(data, dict):
            return DataType.Json
        elif isinstance(data, str):
            return DataType.String
        elif isinstance(data, int):
            return DataType.Int
        elif isinstance(data, float):
            return DataType.Float
        elif isinstance(data, (tuple, list, set)):
            return DataType.List
        return DataType.Unknown


class Constant:
    """
    常量
    """
    
    CATEGORY = "destiny/工具"

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "const_type": ([DataType.String, DataType.Int, DataType.Float], {
                    "default": DataType.String
                }),
                "value": ("STRING", {"default": ""})
            }
        }

    RETURN_TYPES = ("*", )
    RETURN_NAMES = ("output",)

    FUNCTION = "execute"

    def execute(self, const_type, value):
        if const_type == DataType.Int:
            value = int(value)
        elif const_type == DataType.Float:
            value = float(value)
        return (value, )


class Translate:
    """
    翻译
    """
    
    CATEGORY = "destiny/工具"

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "text": ("STRING", {
                    "default": "",
                    "multiline": True
                }),
                "lang_src": (["zh", "en"], {"default": "zh"}),
                "lang_tgt": (["zh", "en"], {"default": "en"}),
                "secret_id": ("STRING", {"default": None}),
                "secret_key": ("STRING", {"default": None}),
            }
        }

    RETURN_TYPES = ("STRING", )
    RETURN_NAMES = ("result", )
    OUTPUT_NODE = True

    FUNCTION = "execute"

    def execute(self, text: str, lang_src, lang_tgt, secret_id=None, secret_key=None):
        assert len(text) > 0, f"[destiny] 请输入文本！"
        assert secret_id is not None and secret_key is not None, f"[destiny] secret_id 和 secret_key 不能为空！"
        text = text.replace('，', ',')
        trans = TencentTranslate(secret_id=secret_id, secret_key=secret_key)
        result = trans.translate(text=text, source=lang_src, target=lang_tgt)
        if result is None:
            result = ""
        else:
            result = result.replace(", ", ",")
        return (result, )


class LabelCombine:

    CATEGORY = "destiny/工具"

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "method": (["默认", "拼接", "或者"], {"default": "默认"})
            },
            "optional": ContainsAnyDict("STRING")
        }
    
    @classmethod
    def VALIDATE_INPUTS(s, input_types):
        return True

    RETURN_TYPES = ("STRING", )
    RETURN_NAMES = ("labels",)
    OUTPUT_NODE = True
    FUNCTION = "execute"

    def execute(self, method: str, **kwargs):
        results = [value for value in kwargs.values()]
        output = ""
        if method == '或者':
            output = "[" + "|".join(results) + "]"
        elif method == '拼接':
            output = " ".join(results)
        else:
            output = ",".join(results)
        
        return (output, )


class CaptionAdd:
    CATEGORY = "destiny/工具"

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "src_dirname": ("STRING", {"placeholder": "Caption directory"}),
                "caption": ("STRING", {}),
                "extension": (['.txt', '.cap', '.caption'], {"default": ".txt"}),
                "method": (["prefix", "suffix"], {"default": "prefix"})
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("dirname",)
    OUTPUT_NODE = True
    FUNCTION = "execute"

    def execute(self, src_dirname: str, caption: str, extension='.txt', method='prefix'):
        assert os.path.exists(src_dirname) and os.path.isdir(src_dirname), f"[destiny] Invalid directory: {src_dirname}"
        assert caption is not None and len(caption) > 0, f"[destiny] Please input the caption!"
        caption = caption.strip().replace('，', ',')
        if not caption.endswith(','):
            caption += ','
        files = os.listdir(src_dirname)
        for file_name in files:
            _, f_ext = os.path.splitext(file_name)
            if not f_ext == extension:
                continue
            filename = os.path.join(src_dirname, file_name)
            with open(filename, mode='r', encoding='utf-8') as fpr:
                data = fpr.read()
            if method == 'prefix':
                data = caption + data
            elif method == 'suffix':
                data += caption
            with open(filename, mode='w+', encoding='utf-8') as fpw:
                fpw.write(data)
        return (src_dirname, )


class CaptionGenerate:
    CATEGORY = "destiny/工具"

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "src_dirname": ("STRING", {"placeholder": "Caption directory"}),
                "caption": ("STRING", {}),
                "extension": (['.txt', '.cap', '.caption'], {"default": ".txt"})
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("dirname",)
    OUTPUT_NODE = True
    FUNCTION = "execute"

    def execute(self, src_dirname: str, caption: str, extension='.txt'):
        assert os.path.exists(src_dirname) and os.path.isdir(src_dirname), f"[destiny] Invalid directory: {src_dirname}"
        assert caption is not None and len(caption) > 0, f"[destiny] Please input the caption!"
        caption = caption.strip().replace('，', ',')
        if caption.endswith(','):
            caption = caption[:-1]
        files = os.listdir(src_dirname)
        for file_name in files:
            f_name, f_ext = os.path.splitext(file_name)
            if f_ext.lower() not in [".png", ".jpg", ".jpeg"]:
                continue
            filename = os.path.join(src_dirname, f"{f_name}{extension}")
            with open(filename, mode='w+', encoding='utf-8') as fpw:
                fpw.write(caption)
        return (src_dirname, )

